import os
import requests
from urllib.parse import urlparse
from huggingface_hub import hf_hub_download, HfApi

token = "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"  # 保持原有的 token

def test_network_connection(endpoint):
    """测试网络连接"""
    try:
        print(f"🌐 测试网络连接: {endpoint}")
        response = requests.get(f"{endpoint}/api/health", timeout=10)
        print(f"✅ {endpoint} 连接正常，状态码: {response.status_code}")
        return True
    except Exception as e:
        print(f"❌ {endpoint} 连接失败: {e}")
        return False

def parse_hf_url(url):
    """解析 Hugging Face URL，提取仓库ID、文件路径等信息"""
    try:
        parsed = urlparse(url)
        parts = parsed.path.strip("/").split("/")

        # 判断是否为 datasets URL
        is_datasets = parts[0] == "datasets"
        start_idx = 1 if is_datasets else 0  # 如果是 datasets，则索引从 1 开始，否则从 0 开始

        # 检查路径是否足够长，至少需要有 user, repo, branch
        if len(parts) > start_idx + 2:
            user = parts[start_idx]
            repo = parts[start_idx + 1]
            branch = parts[start_idx + 3] if len(parts) > start_idx + 3 else None  # branch 索引动态调整
        else:
            print(f"⚠️ URL 中缺少足够的信息: {url}")
            return None, None, None, None, False

        # 获取文件路径
        file_path = ""
        if len(parts) > start_idx + 4:
            file_path = "/".join(parts[start_idx + 4:])

        repo_id = f"{user}/{repo}"
        return repo_id, file_path, branch, repo, is_datasets

    except Exception as e:
        print(f"解析 Hugging Face URL失败: {url}, 错误: {e}")
        return None, None, None, None, False


def ensure_save_path(save_path):
    """确保保存路径不为空且存在"""
    if not save_path.strip():
        print("⚠️ 取消下载：save_path 为空")
        return False
    os.makedirs(save_path, exist_ok=True)
    return True


def get_filename_from_path(path):
    """从文件路径中提取文件名"""
    return os.path.basename(path) if '.' in os.path.basename(path) else None


def download_file_with_hf(url, filename, save_path):
    """下载文件并保存"""
    try:
        print(f"🔍 开始解析URL: {url}")
        repo_id, file_path, branch, repo, is_datasets = parse_hf_url(url)
        print(f"📋 URL解析结果:")
        print(f"   - repo_id: {repo_id}")
        print(f"   - file_path: {file_path}")
        print(f"   - branch: {branch}")
        print(f"   - repo: {repo}")
        print(f"   - is_datasets: {is_datasets}")
        
        if repo_id is None:
            print(f"⚠️ 无法解析 URL: {url}")
            return False, f"无法解析 URL: {url}"
        
        # 创建子目录：按仓库名分类
        target_dir = os.path.join(save_path, repo_id)
        print(f"📁 创建子目录: {target_dir}")
        
        if not ensure_save_path(target_dir):
            return False, "保存路径无效"

        # 如果 filename 为空，从 file_path 获取文件名
        if not filename.strip():
            filename = get_filename_from_path(file_path)
            print(f"📝 从路径提取文件名: {filename}")
            if filename is None:
                print("⚠️ URL 指向文件夹，开始下载文件夹内容")
                download_folder(repo_id, file_path, save_path, branch)
                return True, f"文件夹下载完成: {file_path}"

        # 尝试多个endpoint，提高成功率
        endpoints = [
            "https://hf-mirror.com",   # 国内镜像站（优先）
            "https://huggingface.co",  # 官方站点（备用）
        ]
        
        # 先测试网络连接
        print("🌐 开始网络连接测试...")
        for endpoint in endpoints:
            test_network_connection(endpoint)
        
        for endpoint in endpoints:
            try:
                print(f"🔄 尝试从 {endpoint} 下载...")
                print(f"📥 下载参数:")
                print(f"   - repo_id: {repo_id}")
                print(f"   - filename: {file_path}")
                print(f"   - local_dir: {target_dir}")
                print(f"   - endpoint: {endpoint}")
                
                # 下载文件 - 参数详细说明：
                # repo_id: 模型仓库ID (格式: "用户名/仓库名")
                # filename: 仓库内文件路径
                # token: HuggingFace访问令牌(用于私有仓库)
                # local_dir: 文件直接保存路径(不经过缓存)
                # etag_timeout: 文件校验超时时间(秒)
                # endpoint: 镜像站地址(国内推荐使用hf-mirror.com)
                # resume_download: 启用断点续传(大文件必备)
                # local_dir_use_symlinks:
                #   False(默认): 直接复制文件实体到目标目录(占用更多空间但更可靠)
                #   True: 使用符号链接指向缓存文件(节省空间但依赖缓存位置不变)
                #   建议生产环境设为False确保文件实体存在
                # force_download: 是否强制重新下载(False优先使用缓存)
                downloaded_path = hf_hub_download(
                    repo_id=repo_id,  # 必填: 模型仓库标识
                    filename=file_path,  # 必填: 文件在仓库中的路径
                    token=token,  # 可选: 访问令牌(私有仓库需要)
                    local_dir=target_dir,  # 可选: 直接保存路径
                    etag_timeout=50,  # 推荐值: 文件校验超时
                    endpoint=endpoint,  # 动态选择endpoint
                    resume_download=True,  # 大文件下载必备
                    local_dir_use_symlinks=False,  # 禁用符号链接(确保文件实体存在)
                    force_download=False  # 优先使用缓存(节省流量)
                )
                
                print(f"✅ 下载成功，文件路径: {downloaded_path}")
                
                # 如果需要重命名文件
                if filename.strip():
                    original_path = downloaded_path
                    new_path = os.path.join(target_dir, filename)
                    if original_path != new_path:
                        print(f"🔄 重命名文件: {os.path.basename(original_path)} -> {filename}")
                        os.rename(original_path, new_path)
                        downloaded_path = new_path

                return True, downloaded_path
                
            except Exception as e:
                print(f"⚠️ {endpoint} 下载失败: {e}")
                print(f"🔍 错误类型: {type(e).__name__}")
                print(f"🔍 错误详情: {str(e)}")
                
                # 检查是否文件已经下载成功（网络中断但文件已存在）
                if file_path:
                    expected_file = os.path.join(target_dir, filename) if filename.strip() else os.path.join(target_dir, os.path.basename(file_path))
                    if os.path.exists(expected_file):
                        file_size = os.path.getsize(expected_file)
                        print(f"✅ 检测到文件已存在: {expected_file} (大小: {file_size} 字节)")
                        return True, expected_file
                
                if endpoint == "https://hf-mirror.com":
                    print("🔄 镜像站失败，尝试官方站点...")
                continue
        
        return False, "所有下载源都失败了"

    except Exception as e:
        error_msg = f"下载失败: {filename}，错误: {e}"
        print(error_msg)
        return False, error_msg


def download_folder(repo_id, folder_path, save_path, branch):
    """下载文件夹内容"""
    # 创建子目录：按仓库名分类
    target_dir = os.path.join(save_path, repo_id)
    print(f"📁 创建子目录: {target_dir}")
    
    # 尝试多个endpoint，提高成功率
    endpoints = [
        "https://hf-mirror.com",   # 国内镜像站（优先）
        "https://huggingface.co",  # 官方站点（备用）
    ]
    
    files = []
    hf_api = None
    
    # 获取文件列表
    for endpoint in endpoints:
        try:
            print(f"🔄 尝试从 {endpoint} 获取文件列表...")
            hf_api = HfApi(endpoint=endpoint)
            files = hf_api.list_repo_files(repo_id, revision=branch)
            print(f"✅ 成功获取文件列表，共 {len(files)} 个文件")
            break
        except Exception as e:
            print(f"⚠️ {endpoint} 获取文件列表失败: {e}")
            if endpoint == "https://hf-mirror.com":
                print("🔄 镜像站失败，尝试官方站点...")
            continue
    
    if not files:
        print("❌ 无法获取文件列表，所有endpoint都失败了")
        return

    if not folder_path.strip():
        folder_files = files
    else:
        folder_files = [f for f in files if f.startswith(folder_path)]

    if not ensure_save_path(target_dir):
        return

    print(f"📁 找到 {len(folder_files)} 个文件")
    success_count = 0
    
    for file_path in folder_files:
        try:
            filename = os.path.basename(file_path)
            print(f"📥 下载: {filename}")
            
            # 尝试多个endpoint下载
            for endpoint in endpoints:
                try:
                    downloaded_path = hf_hub_download(
                        repo_id=repo_id,
                        filename=file_path,
                        token=token,
                        local_dir=target_dir,
                        etag_timeout=50,
                        endpoint=endpoint,
                        resume_download=True,
                        local_dir_use_symlinks=False,
                        force_download=False
                    )
                    success_count += 1
                    print(f"✅ {filename} 下载成功")
                    break
                except Exception as e:
                    print(f"⚠️ {endpoint} 下载 {filename} 失败: {e}")
                    if endpoint == "https://hf-mirror.com":
                        print("🔄 镜像站失败，尝试官方站点...")
                    continue
            
        except Exception as e:
            print(f"❌ 下载 {file_path} 失败: {e}")

    print(f"📊 文件夹下载完成: {success_count}/{len(folder_files)} 个文件")

